"""
常用中间件实现示例
"""
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from fastapi import Request, Response
from loguru import logger
import time

# 请求日志中间件
class RequestLogMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        logger.info(f"RequestLogMiddleware:{request.url.path}")
        start = time.time()
        response = await call_next(request)
        process_ms = (time.time() - start) * 1000
        logger.info(f"{request.method} {request.url.path} - {response.status_code} [{process_ms:.2f}ms] {request.client.host}")
        return response

# 统一响应格式中间件（演示，仅对 application/json 有效）
from fastapi.responses import JSONResponse, Response

class WrapResponseMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        # 只处理标准 Response/JSONResponse，有 .body 属性的类型
        if isinstance(response, JSONResponse) or (isinstance(response, Response) and hasattr(response, "body")):
            if 'application/json' in response.headers.get('content-type', '') and hasattr(response, 'body') and response.body:
                import json
                content = json.loads(response.body.decode('utf-8'))
                if isinstance(content, dict) and not {'message', 'data'} <= set(content):
                    content = {'message': 'OK', 'data': content}
                    response.body = json.dumps(content).encode('utf-8')
                    response.headers['content-length'] = str(len(response.body))
        return response
# 全局异常捕获中间件（示例，建议以 FastAPI 原生 Exception Handler 为主）
from starlette.responses import JSONResponse

class ExceptionHandlingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        logger.info(f"ExceptionHandlingMiddleware:{request.url.path}")
        try:
            return await call_next(request)
        except Exception as exc:
            logger.error(f"Exception: {exc}")
            return JSONResponse(status_code=500, content={"message": "Internal server error"})

# 限流（需 slowapi，或参考该实现自定义）
try:
    from slowapi import Limiter, _rate_limit_exceeded_handler
    from slowapi.util import get_remote_address
    from fastapi import FastAPI
    limiter = Limiter(key_func=get_remote_address)
except ImportError:
    limiter = None

# Host 头拦截防护
from fastapi.middleware.trustedhost import TrustedHostMiddleware
# 在 main.py 注册 TrustedHostMiddleware, allowed_hosts=['.xx.cn','localhost']

# 签名校验中间件（伪逻辑流程注释）
class SignatureMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        logger.info(f"SignatureMiddleware: {request.url.path}")
        # 实际应解析 Signature/Header 并验签（略)
        sign = request.headers.get("x-signature")
        if not sign or sign != "secure-token":  # TODO: 接入真正加解密校验
            return JSONResponse(status_code=403, content={"message": "Signature invalid"})
        return await call_next(request)

# 敏感字段自动脱敏中间件（示例: 自动脱敏手机号、邮箱字段）
import re
class FieldDesensitizeMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        logger.info(f"FieldDesensitizeMiddleware: {request.url.path}")

        response = await call_next(request)
        # 仅处理JSON响应且含手机/邮箱字段
        if 'application/json' in response.headers.get('content-type', '') and hasattr(response, 'body'):
            import json
            content = json.loads(response.body.decode('utf-8'))
            def _mask(value):
                if isinstance(value, str):
                    if re.match(r"1[3-9]\\d{9}", value):
                        return value[:3] + '****' + value[-4:]
                    if '@' in value and '.' in value:  # 简单邮箱脱敏
                        idx = value.index('@')
                        return value[0] + '***' + value[idx:]
                return value
            def _walk(d):
                for k, v in d.items():
                    if isinstance(v, dict):
                        d[k] = _walk(v)
                    else:
                        d[k] = _mask(v)
                return d
            if isinstance(content, dict):
                content = _walk(content)
                response.body = json.dumps(content).encode('utf-8')
                response.headers['content-length'] = str(len(response.body))
        return response

# 伪限流中间件（演示，无slowapi时可用）
class FakeRateLimitMiddleware(BaseHTTPMiddleware):
    _ip_counter = {}  # 简单内存IP计数
    WINDOW = 60  # 秒
    LIMIT = 10   # 每窗口最大
    from time import time
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        logger.info(f"FakeRateLimitMiddleware: {request.url.path}")

        now = int(self.time())
        ip = request.client.host
        bucket = self._ip_counter.setdefault(ip, [])
        bucket = [ts for ts in bucket if now-ts < self.WINDOW]
        if len(bucket) >= self.LIMIT:
            return JSONResponse(status_code=429, content={"detail": "Rate limit exceeded"})
        bucket.append(now)
        self._ip_counter[ip] = bucket
        return await call_next(request)
